{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f4abc589d9bfffca",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Usage\n",
    "\n",
    "## 1. Install necessary libs\n",
    "```shell\n",
    "!pip install modelscope\n",
    "!pip install transformers -U\n",
    "!pip install llama-index llama-index-llms-huggingface ipywidgets \n",
    "```\n",
    "\n",
    "## 2. Download data files we need in this example\n",
    "```shell\n",
    "!wget https://modelscope.oss-cn-beijing.aliyuncs.com/resource/rag/punkt.zip\n",
    "!wget https://modelscope.oss-cn-beijing.aliyuncs.com/resource/rag/stopwords.zip\n",
    "!wget https://modelscope.oss-cn-beijing.aliyuncs.com/resource/rag/xianjiaoda.md\n",
    "\n",
    "!mkdir -p /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/tokenizers\n",
    "!mkdir -p /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/corpora\n",
    "\n",
    "!cp /mnt/workspace/punkt.zip /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/tokenizers\n",
    "!cp /mnt/workspace/stopwords.zip /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/corpora\n",
    "!cd /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/tokenizers; unzip punkt.zip;\n",
    "!cd /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/corpora; unzip stopwords.zip;\n",
    "\n",
    "\n",
    "!mkdir -p /mnt/workspace/custom_data\n",
    "!mv /mnt/workspace/xianjiaoda.md /mnt/workspace/custom_data\n",
    "\n",
    "!cd /mnt/workspace\n",
    "```\n",
    "\n",
    "## 3. Go!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c32122833dd7b8c8",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "!pip install modelscope\n",
    "!pip install transformers -U\n",
    "!pip install llama-index llama-index-llms-huggingface ipywidgets "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63704e2b21a9ba52",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "!wget https://modelscope.oss-cn-beijing.aliyuncs.com/resource/rag/punkt.zip\n",
    "!wget https://modelscope.oss-cn-beijing.aliyuncs.com/resource/rag/stopwords.zip\n",
    "!wget https://modelscope.oss-cn-beijing.aliyuncs.com/resource/rag/xianjiaoda.md\n",
    "\n",
    "!mkdir -p /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/tokenizers\n",
    "!mkdir -p /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/corpora\n",
    "\n",
    "!cp /mnt/workspace/punkt.zip /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/tokenizers\n",
    "!cp /mnt/workspace/stopwords.zip /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/corpora\n",
    "!cd /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/tokenizers; unzip punkt.zip;\n",
    "!cd /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/corpora; unzip stopwords.zip;\n",
    "\n",
    "\n",
    "!mkdir -p /mnt/workspace/custom_data\n",
    "!mv /mnt/workspace/xianjiaoda.md /mnt/workspace/custom_data\n",
    "\n",
    "!cd /mnt/workspace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "eef67659e94045c5",
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading Model to directory: /mnt/workspace/.cache/modelscope/qwen/Qwen1.5-4B-Chat\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-01-13 15:52:53,260 - modelscope - INFO - Model revision not specified, using default: [master] version.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-01-13 15:52:53,637 - modelscope - INFO - Creating symbolic link [/mnt/workspace/.cache/modelscope/qwen/Qwen1.5-4B-Chat].\n",
      "2025-01-13 15:52:53,638 - modelscope - WARNING - Failed to create symbolic link /mnt/workspace/.cache/modelscope/qwen/Qwen1.5-4B-Chat for /mnt/workspace/.cache/modelscope/qwen/Qwen1___5-4B-Chat.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4523c5dd31ba411d95cc0ce9e5da8ded",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "llm created\n",
      "Downloading Model to directory: /mnt/workspace/.cache/modelscope/damo/nlp_gte_sentence-embedding_chinese-base\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-01-13 15:53:01,651 - modelscope - INFO - Model revision not specified, using default: [master] version.\n",
      "2025-01-13 15:53:01,894 - modelscope - INFO - initiate model from /mnt/workspace/.cache/modelscope/hub/damo/nlp_gte_sentence-embedding_chinese-base\n",
      "2025-01-13 15:53:01,895 - modelscope - INFO - initiate model from location /mnt/workspace/.cache/modelscope/hub/damo/nlp_gte_sentence-embedding_chinese-base.\n",
      "2025-01-13 15:53:01,898 - modelscope - INFO - initialize model from /mnt/workspace/.cache/modelscope/hub/damo/nlp_gte_sentence-embedding_chinese-base\n",
      "2025-01-13 15:53:02,532 - modelscope - WARNING - No preprocessor field found in cfg.\n",
      "2025-01-13 15:53:02,533 - modelscope - WARNING - No val key and type key found in preprocessor domain of configuration.json file.\n",
      "2025-01-13 15:53:02,533 - modelscope - WARNING - Cannot find available config to build preprocessor at mode inference, current config: {'model_dir': '/mnt/workspace/.cache/modelscope/hub/damo/nlp_gte_sentence-embedding_chinese-base'}. trying to build by task and model information.\n",
      "2025-01-13 15:53:02,588 - modelscope - WARNING - No preprocessor field found in cfg.\n",
      "2025-01-13 15:53:02,588 - modelscope - WARNING - No val key and type key found in preprocessor domain of configuration.json file.\n",
      "2025-01-13 15:53:02,589 - modelscope - WARNING - Cannot find available config to build preprocessor at mode inference, current config: {'model_dir': '/mnt/workspace/.cache/modelscope/hub/damo/nlp_gte_sentence-embedding_chinese-base', 'sequence_length': 128}. trying to build by task and model information.\n",
      "/root/miniconda3/envs/modelscope_1.21/lib/python3.9/site-packages/transformers/modeling_utils.py:1044: FutureWarning: The `device` argument is deprecated and will be removed in v5 of Transformers.\n",
      "  warnings.warn(\n",
      "/root/miniconda3/envs/modelscope_1.21/lib/python3.9/site-packages/transformers/generation/configuration_utils.py:628: UserWarning: `do_sample` is set to `False`. However, `temperature` is set to `0.0` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `temperature`.\n",
      "  warnings.warn(\n",
      "/root/miniconda3/envs/modelscope_1.21/lib/python3.9/site-packages/transformers/generation/configuration_utils.py:633: UserWarning: `do_sample` is set to `False`. However, `top_p` is set to `0.8` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `top_p`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "作为一所历史悠久的综合性研究型大学，西安交通大学有着丰富的校训文化。其中，\"厚德博学，求是创新\"是其最为人所知的校训之一。这句校训不仅体现了学校的教育理念，也反映了学校对学生的期望和要求。此外，西安交通大学还有一句著名的校训：\"明德尚志，自强不息\"，这也是学校对学生的一种激励和引导。这两句校训都强调了教育的重要性，以及学生应该具备的道德品质和自我提升的精神。\n"
     ]
    }
   ],
   "source": [
    "import logging\n",
    "import sys\n",
    "from abc import ABC\n",
    "from typing import Any, List\n",
    "\n",
    "import torch\n",
    "from llama_index.core import (\n",
    "    SimpleDirectoryReader,\n",
    "    VectorStoreIndex,\n",
    "    Settings\n",
    ")\n",
    "from llama_index.core.base.embeddings.base import BaseEmbedding, Embedding\n",
    "from llama_index.core.prompts import PromptTemplate\n",
    "from llama_index.llms.huggingface import HuggingFaceLLM\n",
    "\n",
    "from modelscope import snapshot_download\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
    "\n",
    "# download QWEN model from modelscope\n",
    "qwen15_4B_CHAT = \"qwen/Qwen1.5-4B-Chat\"\n",
    "selected_model = snapshot_download(qwen15_4B_CHAT)\n",
    "\n",
    "# define sys prompt\n",
    "SYSTEM_PROMPT = \"\"\"You are a helpful AI assistant.\"\"\"\n",
    "query_wrapper_prompt = PromptTemplate(\n",
    "    \"[INST]<<SYS>>\\n\" + SYSTEM_PROMPT + \"<</SYS>>\\n\\n{query_str}[/INST] \"\n",
    ")\n",
    "\n",
    "# create HuggingFaceLLM with qwen1.5 \n",
    "llm = HuggingFaceLLM(\n",
    "    context_window=4096,\n",
    "    max_new_tokens=2048,\n",
    "    generate_kwargs={\"temperature\": 0.0, \"do_sample\": False},\n",
    "    query_wrapper_prompt=query_wrapper_prompt,\n",
    "    tokenizer_name=selected_model,\n",
    "    model_name=selected_model,\n",
    "    device_map=\"auto\",\n",
    "    # change these settings below depending on your GPU\n",
    "    model_kwargs={\"torch_dtype\": torch.float16},\n",
    ")\n",
    "print(\"llm created\")\n",
    "\n",
    "\n",
    "# wrap modelscope embedding for llama-index (based on BaseEmbedding)\n",
    "class ModelScopeEmbeddings4LlamaIndex(BaseEmbedding, ABC):\n",
    "    embed: Any = None\n",
    "    model_id: str = \"damo/nlp_gte_sentence-embedding_chinese-base\"\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            model_id: str,\n",
    "            **kwargs: Any,\n",
    "    ) -> None:\n",
    "        super().__init__(**kwargs)\n",
    "        try:\n",
    "            from modelscope.models import Model\n",
    "            from modelscope.pipelines import pipeline\n",
    "            from modelscope.utils.constant import Tasks\n",
    "            # 使用modelscope的embedding模型（包含下载）\n",
    "            self.embed = pipeline(Tasks.sentence_embedding, model=self.model_id)\n",
    "\n",
    "        except ImportError as e:\n",
    "            raise ValueError(\n",
    "                \"Could not import some python packages.\" \"Please install it with `pip install modelscope`.\"\n",
    "            ) from e\n",
    "\n",
    "    def _get_query_embedding(self, query: str) -> Embedding:\n",
    "        text = query.replace(\"\\n\", \" \")\n",
    "        inputs = {\"source_sentence\": [text]}\n",
    "        # note that we have to call tolist() to change numpy.ndarray into python list\n",
    "        return self.embed(input=inputs)['text_embedding'][0].tolist()\n",
    "\n",
    "    def _get_text_embedding(self, text: str) -> Embedding:\n",
    "        text = text.replace(\"\\n\", \" \")\n",
    "        inputs = {\"source_sentence\": [text]}\n",
    "        return self.embed(input=inputs)['text_embedding'][0].tolist()\n",
    "\n",
    "    def _get_text_embeddings(self, texts: List[str]) -> List[Embedding]:\n",
    "        texts = list(map(lambda x: x.replace(\"\\n\", \" \"), texts))\n",
    "        inputs = {\"source_sentence\": texts}\n",
    "        return self.embed(input=inputs)['text_embedding'].tolist()\n",
    "\n",
    "    async def _aget_query_embedding(self, query: str) -> Embedding:\n",
    "        return self._get_query_embedding(query)\n",
    "\n",
    "\n",
    "embedding_model = \"damo/nlp_gte_sentence-embedding_chinese-base\"\n",
    "embeddings = ModelScopeEmbeddings4LlamaIndex(model_id=embedding_model)\n",
    "Settings.embed_model = embeddings\n",
    "Settings.llm = llm\n",
    "\n",
    "# load example documents\n",
    "documents = SimpleDirectoryReader(\"/mnt/workspace/custom_data/\").load_data()\n",
    "\n",
    "# create Vector DB\n",
    "index = VectorStoreIndex.from_documents(documents)\n",
    "\n",
    "# set Logging to DEBUG for more detailed outputs\n",
    "query_engine = index.as_query_engine()\n",
    "\n",
    "# do query\n",
    "response = query_engine.query(\"西安较大的校训是什么\")\n",
    "print(response)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
